#!/bin/bash

PROJECT="pwnker"
AUTHOR="Andrew Bae"
GITHUB="github.com/andrewbae"
DESCRIPTION="A docker environment for pwners"

VERSIONS=("16.04" "17.04" "18.04" "19.04" "20.04")

INSTALL_DIR="/usr/local/bin"
PWNKER_DIR="${INSTALL_DIR}/pwnker-dir"

# Colors preferences
RED="\033[1;31m"
YELLOW="\033[1;33m"
GREEN="\033[1;32m"

NC="\033[0m" # no color

function log() # info, warning, error 
{
	case "$1" in
		info)
			echo -e "${NC}[${GREEN}INFO${NC}] $2"	
			;;
		warning)
			echo -e "${NC}[${YELLOW}WARNING${NC}] $2"	
			;;
		error)
			echo -e "${NC}[${RED}ERROR${NC}] $2"	
			;;
	esac
}

function sanity_check()
{
	if  [ "${#}" != 2 ] || \
		([ "${1}" != "run" ] && [ "${1}" != "build" ] && [ "${1}" != "kill" ]) || \
		([ ! `echo "${2}" | grep -E "([0-9]{2}\.[0-9]{2})"` ] && [ "${2}" != "all" ]); then 
			log error "usage: ${0} [run] [build] [kill] <version>"
			log error "versions: all, `for v0 in ${VERSIONS[@]}; do echo -e $v0", \c"; done`"
			return 1
	fi
	if [ "${1}" == "run" ] && [ "${2}" == "all" ] ; then
		log error "usage: ${0} build all";
		return 1
	fi
	return 0
}

function argument_handling() 
{
	case "${1}" in
		"build")			
			if [ "${2}" == "all" ]; then
				for version in ${VERSIONS[@]}
					do 
						docker build -t "pwnker-ubuntu-${version}" "${PWNKER_DIR}/ubuntu-${version}" --no-cache=false
					done
			else
				docker build -t "pwnker-ubuntu-${2}" "${PWNKER_DIR}/ubuntu-${2}" --no-cache=false
			fi
		;;
		"run")
			if [ "`docker image ls -a | grep "pwnker-ubuntu-${2}"`" ]; then
				if ( ! [ "`docker container ls -a | grep "pwnker-ubuntu-${2}"`" ]); then
					docker run -d --rm -h "pwnker-ubuntu-${2}" --name "pwnker-ubuntu-${2}" -v $HOME/:/pwn --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -t "pwnker-ubuntu-${2}" /bin/bash 
				fi
				current_path=$(echo $PWD | sed "s/`echo $HOME | sed "s/\//\\\\\\\\\//g"`/\//g")
				docker exec -it "pwnker-ubuntu-${2}" /bin/bash -c "cd /pwn/${current_path} && /bin/bash"
			else
				log error "Cannot found the corresponding container. Make sure to build this container before running this container."	
				log warning "Try: pwnker build ${2}"
				return 1
			fi
		;;
		"kill")
			if [ "${2}" == "all" ]; then
				for version in ${VERSIONS[@]}
					do 
						docker kill "pwnker-ubuntu-${version}" 2>/dev/zero 1>/dev/zero
						if [ ${?} == 0 ]; then log info "pwnker-ubuntu${version} is killed";
						else log warning "pwnker-ubuntu${version} is already killed"; fi
					done
			else
				if [ "`docker container ls -a | grep "pwnker-ubuntu-${2}"`" ]; then
						docker kill "pwnker-ubuntu-${2}" 2>/dev/zero 1>/dev/zero
						if [ ${?} == 0 ]; then log info "pwnker-ubuntu${2} is killed";
						else log warning "pwnker-ubuntu${2} is already killed"; fi
				else
					log warning "Cannot found the corresponding container."
					return 1
				fi
			fi			
		;;
	esac
	return 0
}

function main()
{
	log info "${YELLOW}${PROJECT} - ${DESCRIPTION}${NC}"
	log info "Author - ${GITHUB}"
	
	sanity_check $@
	if [ $? == 1 ]; then exit; fi
	
	argument_handling $@
	if [ $? == 1 ]; then exit; fi
}

main $@
